from kaffe.tensorflow import Network
import tensorflow as tf


class VGG16(Network):
    def setup(self):
        (self.feed('data')
             .conv(3, 3, 64, 1, 1, name='conv1_1')
             .conv(3, 3, 64, 1, 1, name='conv1_2')
             .max_pool(2, 2, 2, 2, name='pool1')
             .conv(3, 3, 128, 1, 1, name='conv2_1')
             .conv(3, 3, 128, 1, 1, name='conv2_2')
             .max_pool(2, 2, 2, 2, name='pool2')
             .conv(3, 3, 256, 1, 1, name='conv3_1')
             .conv(3, 3, 256, 1, 1, name='conv3_2')
             .conv(3, 3, 256, 1, 1, name='conv3_3')
             .max_pool(2, 2, 2, 2, name='pool3')
             .conv(3, 3, 512, 1, 1, name='conv4_1')
             .conv(3, 3, 512, 1, 1, name='conv4_2')
             .conv(3, 3, 512, 1, 1, name='conv4_3')
             .max_pool(2, 2, 2, 2, name='pool4')
             .conv(3, 3, 512, 1, 1, name='conv5_1')
             .conv(3, 3, 512, 1, 1, name='conv5_2')
             .conv(3, 3, 512, 1, 1, name='conv5_3')
             .max_pool(2, 2, 2, 2, name='pool5'))
             # .fc(4096, name='fc6')
             # .fc(4096, name='fc7')
             # .fc(1000, relu=False, name='fc8')
             # .softmax(name='prob'))


def vgg_network(input_node, batch_size):
    vgg_net = VGG16({"data": input_node}, trainable=True)
    pool3 = vgg_net.layers["pool3"]
    pool4 = vgg_net.layers["pool4"]
    pool5 = vgg_net.layers["pool5"]
    with tf.variable_scope("pool5_up_sample"):
        weights = tf.get_variable(name="weights", shape=[2, 2, 512, 512], dtype=tf.float32,
                                  initializer=tf.contrib.layers.xavier_initializer_conv2d())
        pool5_up_sample = tf.nn.conv2d_transpose(value=pool5, filter=weights,
                                                 output_shape=[batch_size, 12, 20, 512], strides=[1, 2, 2, 1])
    fcn_16s = tf.add(pool4, pool5_up_sample)
    with tf.variable_scope("pool4_up_sample"):
        weights = tf.get_variable(name="weights", shape=[2, 2, 256, 512], dtype=tf.float32,
                                  initializer=tf.contrib.layers.xavier_initializer_conv2d())
        pool4_up_sample = tf.nn.conv2d_transpose(value=fcn_16s, filter=weights,
                                                 output_shape=[batch_size, 23, 40, 256], strides=[1, 2, 2, 1])
    fcn_8s = tf.add(pool3, pool4_up_sample)
    with tf.variable_scope("vgg_output"):
        weights = tf.get_variable(name="weights", shape=[8, 8, 1, 256], dtype=tf.float32,
                                  initializer=tf.contrib.layers.xavier_initializer_conv2d())
        output = tf.nn.conv2d_transpose(value=fcn_8s, filter=weights,
                                        output_shape=[batch_size, 180, 320, 1], strides=[1, 8, 8, 1])
    return vgg_net, output


def main():
    input_node = tf.placeholder(dtype=tf.float32, shape=[None, 180, 320, 3])
    vgg_net, result = vgg_network(input_node, batch_size=32)
    print(result.get_shape())
    # with tf.Session() as session:
    #     session.run(tf.initialize_all_variables())
    #     vgg_net.load("./VGG16.npy", session, ignore_missing=True)


if __name__ == "__main__":
    main()
